import os
import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# 创建必要的目录结构
data_dir = 'd:/.homework/artificial_intelligence/work8/data'
output_dir = 'd:/.homework/artificial_intelligence/work8/output'

# 确保目录存在
os.makedirs(data_dir, exist_ok=True)
os.makedirs(output_dir, exist_ok=True)

def load_data(batch_size=64):
    """
    加载MNIST数据集
    
    参数:
        batch_size: 批次大小
        
    返回:
        train_loader: 训练数据加载器
        test_loader: 测试数据加载器
    """
    # 定义数据转换：将图像转换为张量并进行归一化
    transform = transforms.Compose([
        transforms.Resize((32, 32)),  # LeNet-5期望32x32输入
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))  # MNIST数据集的均值和标准差
    ])
    
    # 加载训练集和测试集
    train_dataset = datasets.MNIST(data_dir, train=True, download=True, transform=transform)
    test_dataset = datasets.MNIST(data_dir, train=False, download=True, transform=transform)
    
    # 创建数据加载器
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
    
    return train_loader, test_loader

def get_paths():
    """返回数据和输出目录路径"""
    return data_dir, output_dir